# Copyright (c) 2018 Presto Labs Pte. Ltd.
# Author: leon

import os

from absl import app, flags

import pandas
import matplotlib
matplotlib.use('Agg')
import matplotlib.pyplot as plt

FLAGS = flags.FLAGS


def do_plot_sim(df_edge_2, df_edge_4, df_edge_6, csv_root):
  plt.figure(figsize=(10, 10))
  plt.title('Bitflyer -> Binance BTC-JPY unhedge mm simulation result')
  plt.xticks(df_edge_2.index, df_edge_2['ma_window'])
  edge_2, = plt.plot(df_edge_2.index, df_edge_2['pnl'], color='g', marker='o')
  edge_4, = plt.plot(df_edge_2.index, df_edge_4['pnl'], color='b', marker='o')
  edge_6, = plt.plot(df_edge_2.index, df_edge_6['pnl'], color='r', marker='o')
  plt.xlabel('ma_window')
  plt.ylabel('pnl JPY')
  plt.grid(True)
  plt.legend([edge_2, edge_4, edge_6], ['edge_bp=2', 'edge_bp=4', 'edge_bp=6'])
  plt.savefig(os.path.join(csv_root, 'Bitflyer_BTC-JPY_Binance_sim_result.png'))


def plot_sim(df, csv_root):
  df_edge_2 = df.loc[df['edge_bp'] == 2]
  df_edge_4 = df.loc[df['edge_bp'] == 4]
  df_edge_6 = df.loc[df['edge_bp'] == 6]
  do_plot_sim(df_edge_2, df_edge_4, df_edge_6, csv_root)


def read_csv_into_df(csv_root, csv):
  csv_dir = os.path.join(csv_root, csv)
  df = pandas.read_csv(csv_dir, sep=',', header=0)
  return df


def add_columns_and_merge(df):
  df['ma_window'] = df['name'].str.split('.').str[0].str[:2]
  df['edge_bp'] = df['name'].str.split('.').str[1].str[:2]
  df['trading_date'] = df['name'].str.split('.').str[-1]
  df.drop(['name'], inplace=True, axis=1)
  df[['ma_window', 'edge_bp']] = df[['ma_window', 'edge_bp']].astype(float)
  df[['pnl', 'pq_sum']] = df[['pnl', 'pq_sum']].astype(float)
  df = df.groupby(['ma_window', 'edge_bp']).agg('mean').reset_index()
  df.rename(columns={'pq_sum': 'pq'}, inplace=True)
  return df


def read_tabulate_into_df(tabulate_file_path):
  line_index = 0
  headers = None
  rows = []
  with open(tabulate_file_path, 'r') as infile:
    for line in infile:
      if line_index == 0:
        headers = line.strip().split()
      elif line_index == 1:
        pass
      else:
        row = line.strip().split()
        rows.append(row)
      line_index += 1
  df = pandas.DataFrame(rows, columns=headers)
  return df


def main(argv):
  csv_root = FLAGS.csv_root
  tabulate_file_path = FLAGS.tabulate_file_path
  assert csv_root or tabulate_file_path, '--csv_root must be specified.'

  if tabulate_file_path is not None:
    sim_df = read_tabulate_into_df(tabulate_file_path)
    sim_df = add_columns_and_merge(sim_df)
    csv_root = ''
  else:
    sim_df = read_csv_into_df(csv_root, 'sim_all.csv.binance')
    sim_df = add_columns_and_merge(sim_df)

  plot_sim(sim_df, csv_root)

  return 0


if __name__ == '__main__':
  flags.DEFINE_string('csv_root', None, 'Input csv files root directory.')

  flags.DEFINE_string('tabulate_file_path', None, 'Input tabulate file path.')

  app.run(main)
